#!/bin/bash

# ABACUS势能面扫描脚本 - 三角形区域 + Z方向扫描（顺序执行）
# 包含势能面绘图功能

# 设置ABACUS执行路径
ABACUS_PATH="/usr/local/bin/abacus"
CORE=10  # 设置MPI进程数

# 三角形顶点坐标（分数坐标）
A=(0.0 0.0 0.09)
B=(0.25 0.25 0.09)
C=(0.5 0.0 0.09)

# Z方向扫描范围
Z_START=0.04
Z_END=0.24
Z_STEP=0.01  # Z方向步长（分数坐标）

# 扫描步长（埃）
STEP_XY=0.4

# 晶格常数（用于埃到分数坐标的转换）- 根据新STRU文件更新
LATTICE_A=12.68000000000
LATTICE_B=10.98120211998

# 计算XY步长对应的分数坐标
STEP_X=$(echo "scale=6; $STEP_XY / $LATTICE_A" | bc -l)
STEP_Y=$(echo "scale=6; $STEP_XY / $LATTICE_B" | bc -l)

echo "XY步长转换: $STEP_XY 埃 -> X方向: $STEP_X, Y方向: $STEP_Y 分数坐标"
echo "Z扫描范围: $Z_START 到 $Z_END, 步长: $Z_STEP"

# 创建工作目录
WORK_DIR="PES_scan_3D_sequential"
mkdir -p $WORK_DIR
cd $WORK_DIR

# 复制输入文件
cp ../INPUT ./
cp ../STRU ./


# 计数器
total_count=0

# 判断点是否在三角形内的函数
point_in_triangle() {
    local x=$1
    local y=$2
    
    # 计算重心坐标
    local denom=$(echo "(${B[1]} - ${C[1]}) * (${A[0]} - ${C[0]}) + (${C[0]} - ${B[0]}) * (${A[1]} - ${C[1]})" | bc -l)
    if [ $(echo "$denom == 0" | bc -l) -eq 1 ]; then
        echo "0"
        return
    fi
    
    local alpha=$(echo "((${B[1]} - ${C[1]}) * ($x - ${C[0]}) + (${C[0]} - ${B[0]}) * ($y - ${C[1]})) / $denom" | bc -l)
    local beta=$(echo "((${C[1]} - ${A[1]}) * ($x - ${C[0]}) + (${A[0]} - ${C[0]}) * ($y - ${C[1]})) / $denom" | bc -l)
    local gamma=$(echo "1 - $alpha - $beta" | bc -l)
    
    # 判断点是否在三角形内（包括边界）
    local in_triangle=$(echo "$alpha >= -0.0001 && $beta >= -0.0001 && $gamma >= -0.0001" | bc -l)
    echo $in_triangle
}

echo "开始3D势能面扫描（三角形XY区域 + Z方向）..."

# 在三角形区域内扫描XY坐标，对每个XY点扫描Z方向
for i in $(seq 0 0.05 0.5); do
    for j in $(seq 0 0.05 0.5); do
        x=$(echo "$i" | bc -l)
        y=$(echo "$j" | bc -l)
        
        # 检查点是否在三角形内
        in_triangle=$(point_in_triangle $x $y)
        
        if [ $in_triangle -eq 1 ]; then
            # 对每个XY点，扫描Z方向
            for z in $(seq $Z_START $Z_STEP $Z_END); do
                # 创建计算目录
                x_formatted=$(printf "%.4f" $x)
                y_formatted=$(printf "%.4f" $y)
                z_formatted=$(printf "%.3f" $z)
                dir_name="scan_${total_count}_x${x_formatted}_y${y_formatted}_z${z_formatted}"
                mkdir -p $dir_name
                cd $dir_name
                
                # 复制INPUT文件
                cp ../INPUT ./
                
                # 生成新的STRU文件 - 使用您提供的新结构
                cat > STRU << EOF
ATOMIC_SPECIES
C   12.011  C_ONCV_PBE-1.0.upf
Al   26.982  Al_ONCV_PBE-1.0.upf

NUMERICAL_ORBITAL
C_gga_10au_100Ry_3s3p2d.orb
Al_gga_10au_100Ry_6s6p2d.orb

LATTICE_CONSTANT
1.889726

LATTICE_VECTORS
12.68000000000  0.000000000000  0.000000000000
-6.34000000000  10.98120211998  0.000000000000
0.000000000000  0.000000000000  20.00000000000

ATOMIC_POSITIONS
Direct

C 
0.000
48
0.129830000000  0.064920000000  0.000000000000  1  1  1  mag  0.0
0.435090000000  0.064910000000  0.000000000000  1  1  1  mag  0.0
0.435090000000  0.370170000000  0.000000000000  1  1  1  mag  0.0
0.370170000000  0.435090000000  0.000000000000  1  1  1  mag  0.0
0.064920000000  0.435090000000  0.000000000000  1  1  1  mag  0.0
0.064920000000  0.129830000000  0.000000000000  1  1  1  mag  0.0
0.258770000000  0.129390000000  0.000000000000  1  1  1  mag  0.0
0.370620000000  0.129380000000  0.000000000000  1  1  1  mag  0.0
0.369800000000  0.240410000000  0.000000000000  1  1  1  mag  0.0
0.241240000000  0.370620000000  0.000000000000  1  1  1  mag  0.0
0.129390000000  0.370620000000  0.000000000000  1  1  1  mag  0.0
0.129380000000  0.258770000000  0.000000000000  1  1  1  mag  0.0
0.629830000000  0.064920000000  0.000000000000  1  1  1  mag  0.0
0.935090000000  0.064910000000  0.000000000000  1  1  1  mag  0.0
0.935090000000  0.370170000000  0.000000000000  1  1  1  mag  0.0
0.870170000000  0.435090000000  0.000000000000  1  1  1  mag  0.0
0.564920000000  0.435090000000  0.000000000000  1  1  1  mag  0.0
0.564920000000  0.129830000000  0.000000000000  1  1  1  mag  0.0
0.758770000000  0.129390000000  0.000000000000  1  1  1  mag  0.0
0.870620000000  0.129380000000  0.000000000000  1  1  1  mag  0.0
0.869800000000  0.240410000000  0.000000000000  1  1  1  mag  0.0
0.741240000000  0.370620000000  0.000000000000  1  1  1  mag  0.0
0.629390000000  0.370620000000  0.000000000000  1  1  1  mag  0.0
0.629380000000  0.258770000000  0.000000000000  1  1  1  mag  0.0
0.129830000000  0.564920000000  0.000000000000  1  1  1  mag  0.0
0.435090000000  0.564920000000  0.000000000000  1  1  1  mag  0.0
0.435090000000  0.870170000000  0.000000000000  1  1  1  mag  0.0
0.370170000000  0.935090000000  0.000000000000  1  1  1  mag  0.0
0.064920000000  0.935090000000  0.000000000000  1  1  1  mag  0.0
0.064920000000  0.629830000000  0.000000000000  1  1  1  mag  0.0
0.258770000000  0.629390000000  0.000000000000  1  1  1  mag  0.0
0.370620000000  0.629380000000  0.000000000000  1  1  1  mag  0.0
0.369800000000  0.740410000000  0.000000000000  1  1  1  mag  0.0
0.241240000000  0.870620000000  0.000000000000  1  1  1  mag  0.0
0.129390000000  0.870620000000  0.000000000000  1  1  1  mag  0.0
0.129380000000  0.758770000000  0.000000000000  1  1  1  mag  0.0
0.629830000000  0.564920000000  0.000000000000  1  1  1  mag  0.0
0.935090000000  0.564920000000  0.000000000000  1  1  1  mag  0.0
0.935090000000  0.870170000000  0.000000000000  1  1  1  mag  0.0
0.870170000000  0.935090000000  0.000000000000  1  1  1  mag  0.0
0.564920000000  0.935090000000  0.000000000000  1  1  1  mag  0.0
0.564920000000  0.629830000000  0.000000000000  1  1  1  mag  0.0
0.758770000000  0.629390000000  0.000000000000  1  1  1  mag  0.0
0.870620000000  0.629380000000  0.000000000000  1  1  1  mag  0.0
0.869800000000  0.740410000000  0.000000000000  1  1  1  mag  0.0
0.741240000000  0.870620000000  0.000000000000  1  1  1  mag  0.0
0.629390000000  0.870620000000  0.000000000000  1  1  1  mag  0.0
0.629380000000  0.758770000000  0.000000000000  1  1  1  mag  0.0

Al
0.000
1
$x  $y  $z  1  1  1  mag  0.0
EOF
                
                echo "运行计算: $dir_name, 坐标: ($x, $y, $z)"
                
                # 运行ABACUS计算并等待完成
                OMP_NUM_THREADS=1 mpirun -n 10 /usr/local/bin/abacus > out.log 2> err.log &
                pid=$!
                wait $pid
                
                # 检查计算是否成功完成
                if [ -f "OUT.ABACUS/running_scf.log" ] && grep -q "FINAL_ETOT_IS" OUT.ABACUS/running_scf.log; then
                    energy=$(grep "FINAL_ETOT_IS" OUT.ABACUS/running_scf.log | tail -1 | awk '{print $2}')
                    echo "计算成功完成: $dir_name, 能量: $energy Ry"
                else
                    echo "警告: 计算可能未完成: $dir_name"
                fi
                
                cd ..
                ((total_count++))
            done
        fi
    done
done

echo "3D势能面扫描完成！总共计算了 $total_count 个点。"

# 生成结果汇总文件
echo "生成结果汇总..."
echo "x_coord,y_coord,z_coord,energy,force_x,force_y,force_z,stress_xx,stress_yy,stress_zz" > results_3d_scan.csv

for dir in scan_*/; do
    if [ -d "$dir" ]; then
        cd "$dir"
        
        # 从文件名提取坐标
        dir_name=$(basename $dir)
        x_coord=$(echo $dir_name | grep -o 'x[0-9.]*' | sed 's/x//')
        y_coord=$(echo $dir_name | grep -o 'y[0-9.]*' | sed 's/y//')
        z_coord=$(echo $dir_name | grep -o 'z[0-9.]*' | sed 's/z//')
        
        # 提取能量
        energy="NaN"
        if [ -f "OUT.ABACUS/running_scf.log" ]; then
            energy_line=$(grep "FINAL_ETOT_IS" OUT.ABACUS/running_scf.log | tail -1)
            if [ ! -z "$energy_line" ]; then
                energy=$(echo $energy_line | awk '{print $2}')
            fi
        fi
        
        # 提取力
        force_x="NaN"; force_y="NaN"; force_z="NaN"
        if [ -f "OUT.ABACUS/FORCE.txt" ]; then
            # 提取Al原子的力（最后一个原子）
            force_line=$(tail -1 OUT.ABACUS/FORCE.txt)
            if [ ! -z "$force_line" ]; then
                force_x=$(echo $force_line | awk '{print $2}')
                force_y=$(echo $force_line | awk '{print $3}')
                force_z=$(echo $force_line | awk '{print $4}')
            fi
        fi
        
        # 提取应力
        stress_xx="NaN"; stress_yy="NaN"; stress_zz="NaN"
        if [ -f "OUT.ABACUS/STRESS.txt" ]; then
            # 提取应力张量的对角线元素
            stress_line=$(grep "TOTAL STRESS" OUT.ABACUS/STRESS.txt -A 3 | tail -3 | head -1)
            if [ ! -z "$stress_line" ]; then
                stress_xx=$(echo $stress_line | awk '{print $1}')
                stress_yy=$(echo $stress_line | awk '{print $2}')
                stress_zz=$(echo $stress_line | awk '{print $3}')
            fi
        fi
        
        cd ..
        echo "$x_coord,$y_coord,$z_coord,$energy,$force_x,$force_y,$force_z,$stress_xx,$stress_yy,$stress_zz" >> results_3d_scan.csv
    fi
done

echo "3D扫描结果已保存到 results_3d_scan.csv"

# 生成进度统计
success_count=$(grep -c -v "NaN" results_3d_scan.csv)
((success_count--))  # 减去标题行
echo "成功计算的点数: $success_count/$total_count"

# 生成势能面绘图脚本
echo "生成势能面绘图..."
cat > plot_pes.py << 'EOF'
#!/usr/bin/env python3
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
from scipy.interpolate import griddata
from mpl_toolkits.mplot3d import Axes3D
import os

def read_and_process_data():
    """读取和处理数据"""
    try:
        data = pd.read_csv('results_3d_scan.csv')
        print(f"成功读取 {len(data)} 行数据")
        
        # 过滤有效数据
        valid_data = data[data['energy'] != 'NaN'].copy()
        valid_data = valid_data.apply(pd.to_numeric, errors='coerce')
        valid_data = valid_data.dropna(subset=['x_coord', 'y_coord', 'z_coord', 'energy'])
        
        print(f"有效数据点: {len(valid_data)}")
        return valid_data
    except Exception as e:
        print(f"读取数据时出错: {e}")
        return None

def create_2d_contour_plots(valid_data):
    """创建2D等高线图"""
    if len(valid_data) < 10:
        print("数据点不足，无法生成等高线图")
        return
    
    # 为每个Z值创建2D势能面
    z_values = valid_data['z_coord'].unique()
    z_values = np.sort(z_values)
    
    print(f"找到 {len(z_values)} 个不同的Z值")
    
    for i, z_val in enumerate(z_values[:9]):  # 只绘制前9个Z值
        z_data = valid_data[valid_data['z_coord'] == z_val]
        
        if len(z_data) < 5:
            continue
            
        # 创建网格数据
        xi = np.linspace(z_data['x_coord'].min(), z_data['x_coord'].max(), 100)
        yi = np.linspace(z_data['y_coord'].min(), z_data['y_coord'].max(), 100)
        XI, YI = np.meshgrid(xi, yi)
        
        # 插值
        points = z_data[['x_coord', 'y_coord']].values
        values = z_data['energy'].values
        
        try:
            ZI = griddata(points, values, (XI, YI), method='cubic')
        except:
            ZI = griddata(points, values, (XI, YI), method='linear')
        
        # 创建图形
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(16, 6))
        
        # 等高线图
        contour = ax1.contourf(XI, YI, ZI, levels=50, cmap='viridis')
        ax1.scatter(points[:, 0], points[:, 1], c='red', s=20, alpha=0.6, label='计算点')
        ax1.set_xlabel('X坐标')
        ax1.set_ylabel('Y坐标')
        ax1.set_title(f'势能面等高线 (Z={z_val:.3f})')
        plt.colorbar(contour, ax=ax1, label='能量 (Ry)')
        ax1.legend()
        
        # 3D表面图
        ax2 = fig.add_subplot(122, projection='3d')
        surf = ax2.plot_surface(XI, YI, ZI, cmap='viridis', alpha=0.8)
        ax2.scatter(points[:, 0], points[:, 1], values, c='red', s=30, alpha=1.0)
        ax2.set_xlabel('X坐标')
        ax2.set_ylabel('Y坐标')
        ax2.set_zlabel('能量 (Ry)')
        ax2.set_title(f'3D势能面 (Z={z_val:.3f})')
        
        plt.tight_layout()
        plt.savefig(f'pes_contour_z_{z_val:.3f}.png', dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"生成势能面图: pes_contour_z_{z_val:.3f}.png")

def create_energy_vs_z_plot(valid_data):
    """创建能量随Z坐标变化图"""
    if len(valid_data) == 0:
        return
    
    # 找到能量最低的XY点
    min_energy_idx = valid_data['energy'].idxmin()
    best_xy = valid_data.loc[min_energy_idx, ['x_coord', 'y_coord']]
    
    # 提取该XY点的所有Z扫描数据
    xy_data = valid_data[
        (np.abs(valid_data['x_coord'] - best_xy['x_coord']) < 0.001) & 
        (np.abs(valid_data['y_coord'] - best_xy['y_coord']) < 0.001)
    ]
    
    if len(xy_data) > 1:
        plt.figure(figsize=(10, 6))
        plt.plot(xy_data['z_coord'], xy_data['energy'], 'bo-', linewidth=2, markersize=6)
        plt.xlabel('Z坐标')
        plt.ylabel('能量 (Ry)')
        plt.title(f'能量随Z坐标变化 (X={best_xy["x_coord"]:.3f}, Y={best_xy["y_coord"]:.3f})')
        plt.grid(True, alpha=0.3)
        plt.savefig('energy_vs_z.png', dpi=300, bbox_inches='tight')
        plt.close()
        print("生成能量随Z坐标变化图: energy_vs_z.png")

def create_summary_plot(valid_data):
    """创建汇总图"""
    if len(valid_data) < 10:
        return
    
    fig = plt.figure(figsize=(15, 10))
    
    # 1. 最低能量点分布
    ax1 = fig.add_subplot(221)
    scatter = ax1.scatter(valid_data['x_coord'], valid_data['y_coord'], 
                         c=valid_data['energy'], cmap='viridis', s=50)
    ax1.set_xlabel('X坐标')
    ax1.set_ylabel('Y坐标')
    ax1.set_title('势能面扫描点分布')
    plt.colorbar(scatter, ax=ax1, label='能量 (Ry)')
    
    # 2. 能量分布直方图
    ax2 = fig.add_subplot(222)
    ax2.hist(valid_data['energy'], bins=20, alpha=0.7, color='skyblue')
    ax2.set_xlabel('能量 (Ry)')
    ax2.set_ylabel('频数')
    ax2.set_title('能量分布直方图')
    ax2.grid(True, alpha=0.3)
    
    # 3. Z方向力分布
    ax3 = fig.add_subplot(223)
    ax3.scatter(valid_data['z_coord'], valid_data['force_z'], alpha=0.6)
    ax3.set_xlabel('Z坐标')
    ax3.set_ylabel('Z方向力 (Ry/Bohr)')
    ax3.set_title('Z方向力分布')
    ax3.grid(True, alpha=0.3)
    
    # 4. 最低能量路径
    ax4 = fig.add_subplot(224, projection='3d')
    
    # 对每个Z值找到最低能量点
    z_unique = valid_data['z_coord'].unique()
    min_energy_path = []
    for z in z_unique:
        z_subset = valid_data[valid_data['z_coord'] == z]
        if len(z_subset) > 0:
            min_idx = z_subset['energy'].idxmin()
            min_energy_path.append(z_subset.loc[min_idx])
    
    if min_energy_path:
        min_df = pd.DataFrame(min_energy_path)
        scatter = ax4.scatter(min_df['x_coord'], min_df['y_coord'], min_df['z_coord'], 
                             c=min_df['energy'], cmap='viridis', s=50)
        ax4.set_xlabel('X坐标')
        ax4.set_ylabel('Y坐标')
        ax4.set_zlabel('Z坐标')
        ax4.set_title('最低能量路径')
        plt.colorbar(scatter, ax=ax4, label='能量 (Ry)')
    
    plt.tight_layout()
    plt.savefig('pes_summary.png', dpi=300, bbox_inches='tight')
    plt.close()
    print("生成汇总图: pes_summary.png")

def main():
    """主函数"""
    print("开始生成势能面图...")
    
    valid_data = read_and_process_data()
    if valid_data is None or len(valid_data) < 5:
        print("数据不足，无法生成图形")
        return
    
    # 创建输出目录
    os.makedirs('plots', exist_ok=True)
    os.chdir('plots')
    
    # 生成各种图形
    create_2d_contour_plots(valid_data)
    create_energy_vs_z_plot(valid_data)
    create_summary_plot(valid_data)
    
    print("势能面图生成完成！")

if __name__ == "__main__":
    main()
EOF

# 运行绘图脚本
echo "运行势能面绘图..."
python3 plot_pes.py

echo "所有任务完成！"
echo "结果文件: results_3d_scan.csv"
echo "图形文件: 在 plots/ 目录中"